import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class ECALayer(nn.Module):
    '''
    def __init__(self, channel, k_size=3):
        """
        ECA模块：不使用全连接而是使用1D卷积进行通道注意力建模
        :param channel: 输入通道数
        :param k_size: 卷积核大小（通常是奇数：3, 5, 7）
        """
        super(ECALayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 全局平均池化
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()
    '''
    def __init__(self, channel, gamma=2, b=1):
        super(ECALayer, self).__init__()
        # 自适应卷积核大小
        t = int(abs((math.log2(channel) / gamma) + b))
        k_size = t if t % 2 else t + 1

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x):
        y = self.avg_pool(x)  # [B, C, 1, 1]
        y = y.squeeze(-1).transpose(-1, -2)  # [B, 1, C]
        y = self.conv(y)  # [B, 1, C]
        y = self.sigmoid(y).transpose(-1, -2).unsqueeze(-1)  # [B, C, 1, 1]
        return x * y.expand_as(x)